from scipy.integrate import odeint
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import fsolve


plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
plt.rc('font', size=10)
plt.rc('font', family='SimHei')


v = 1
b = 0.55 / (2 * np.pi)
theta_0 = 16 * 2 * np.pi
l_head = 3.41 - 0.275 * 2
l_body = 2.2 - 0.275 * 2
label_list = [0, 1, 51, 101, 151, 201, 222]


# theta与t的关系
def pfun(y, t):
    global v, b
    return (-1) * v / (b * ((1 + y ** 2) ** 0.5))


def x(theta, r):
    return r * np.cos(theta)


def y(theta, r):
    return r * np.sin(theta)


"""
    求解极坐标系下圆和等距螺线的交点

    参数:
        r0: 圆心的极径
        phi0: 圆心的极角(弧度)
        R: 圆的半径
        b: 螺线方程 r = b*theta 中的参数
        theta_guesses: 初始猜测的角度列表(弧度)

    返回:
        交点列表，每个交点格式为 (theta, r)
"""
def find_intersections(r0, phi0, R, b, theta_guesses):
    # 定义方程组
    def equations(p):
        theta = p[0]
        r = b * theta
        eq1 = r**2 + r0**2 - 2*r*r0*np.cos(theta - phi0) - R**2
        return [eq1]

    # 求解每个初始猜测
    solutions = []
    for guess in theta_guesses:
        theta_sol = fsolve(equations, guess)[0]
        r_sol = b * theta_sol
        # 检查解是否合理(避免数值误差导致的无效解)
        if theta_sol > 0 and abs(equations([theta_sol])[0]) < 1e-6:
            # 检查是否已经找到过这个解(避免重复)
            is_new = True
            for sol in solutions:
                if np.allclose([theta_sol], [sol[0]], atol=1e-3):
                    is_new = False
                    break
            if is_new:
                solutions.append((theta_sol, r_sol))

    return solutions


t = np.arange(0, 300, 0.1)  # 创建自变量序列
theta = np.arange(0, 15, 0.1)
soli = np.array(odeint(pfun, theta_0, t)[:, 0])  # 求数值解
print(soli)
# ax = plt.subplot(111, projection='polar')
# ax.plot(theta, b * theta, linewidth=3, color='blue')
# ax.plot(soli, b * soli, linewidth=3, color='red')  # 第一个参数为角度，第二个参数为极径

# 初始猜测角度(需要根据具体问题调整)
def find_next(t, theta=None, r=None, i=0):
    if not theta:
        theta = soli[t * 10]
    if not r:
        r = b * soli[t * 10]
    if i in label_list:
        print(x(theta, r), y(theta, r))
    theta_guesses = [theta + 0.1]  # 初始猜测角度。略大于原点，找到的解也会更大

    # 求解交点
    if i == 0:
        intersections = find_intersections(r, theta, l_head, b, theta_guesses)
    else:
        intersections = find_intersections(r, theta, l_body, b, theta_guesses)

    # # 打印结果
    # print("找到的交点(θ, r):")
    # for i, (theta_result, r_result) in enumerate(intersections):
    #     print(f"交点 {i+1}: θ = {theta_result:.4f} rad, r = {r_result:.4f}")
    #     # 转换为笛卡尔坐标
    #     x_result = r_result * np.cos(theta_result)
    #     y_result = r_result * np.sin(theta_result)
    #     print(f"笛卡尔坐标: ({x_result:.4f}, {y_result:.4f})")


    # plt.plot(x(theta, r), y(theta, r), "ro", label=f'待求圆心')

    # # 绘制圆
    # theta_circle = np.linspace(0, 2 * np.pi, 100)
    # x_circle = r * np.cos(theta) + l_head * np.cos(theta_circle)
    # y_circle = r * np.sin(theta) + l_head * np.sin(theta_circle)
    # plt.plot(x_circle, y_circle, label=f'圆: 中心({r * np.cos(theta):.1f},{r * np.sin(theta):.1f}), 半径{l_head}')

    # 绘制交点
    for theta_result, r_result in intersections:
        x_result = r_result * np.cos(theta_result)
        y_result = r_result * np.sin(theta_result)
        plt.plot(x_result, y_result, 'ro')
        # plt.text(x_result, y_result, f'  ({x_result:.2f}, {y_result:.2f})')
        if i < 223 and theta_result < soli[0]:
            find_next(t, theta_result, r_result, i + 1)


# 绘制图形
plt.figure(figsize=(8, 8))
# 绘制螺线
plt.plot(x(soli, b * soli), y(soli, b * soli), label=f'螺线: r = {b}θ')
find_next(60)
plt.xlabel('X')
plt.ylabel('Y')
plt.title('极坐标系下圆和等距螺线的交点')
plt.axis('equal')
plt.grid(True)
plt.legend()
plt.show()
